clear;

A = [
    0.81 0.13; 
    0.19 0.87
];

B = zeros(2,2,2);
% pi = 1: stay in the state
B(:,:,1) = [
    0.91 0.27; 
    0.09 0.73
];
% pi = 2: translation
B(:,:,2) = [
    0.8 0.44; 
    0.2 0.56
];

D = [
    0.99; 
    0.01
];

E = [
    0.5; 
    0.5
];


%P(o1,o2,s1,s2,pi)
Joint = zeros(2,2,2,2,2);
for o1=1:2
    for o2=1:2
        for s1 = 1:2
            for s2 = 1:2
                for pi = 1:2            
                    Joint(s1,s2,o1,o2,pi) = E(pi) * D(s1) * A(o1,s1) * B(s2,s1,pi) * A(o2,s2);
                end
            end
        end
    end
end

%p(o1)
Po1 = zeros(2,1);
for o1=1:2
        Po1(o1)=sum(Joint(:,:,o1,:,:), "all");
end

% P(o2)
Po2 = zeros(2,1);
for o2=1:2
        Po2(o2)=sum(Joint(:,:,:,o2,:), "all");
end
disp(Po2);

%p(s1,s2,o2,pi|o1)
P_s1_s2_o2_pi_cnd_o1 = zeros(2,2,2,2,2);
for o1=1:2
    for s1=1:2
        for s2=1:2
            for o2=1:2
                for pi=1:2
                    P_s1_s2_o2_pi_cnd_o1(s1,s2,o2,pi,o1) = Joint(s1,s2,o1,o2,pi) / Po1(o1);
                end
            end
        end
    end
end

P_s1_s2_o1_pi_cnd_o2 = zeros(2,2,2,2,2);
for s1=1:2
    for s2=1:2
        for o1=1:2
            for o2=1:2
                for pi=1:2
                    P_s1_s2_o1_pi_cnd_o2(s1,s2,o1,pi,o2) = Joint(s1,s2,o1,o2,pi) / Po2(o2);
                end
            end
        end
    end
end

%p(pi|o1)
P_pi_o1 = zeros(2,1);
for o1=1:2
    for pi=1:2
        P_pi_o1(pi,o1)=sum(P_s1_s2_o2_pi_cnd_o1(:,:,:,pi,o1), "all");
    end
end
disp(P_pi_o1)

%p(pi|o2)
P_pi_o2 = zeros(2,1);
for o2=1:2
    for pi=1:2
        P_pi_o2(pi,o2)=sum(P_s1_s2_o1_pi_cnd_o2(:,:,:,pi,o2), "all");
    end
end
disp(P_pi_o2)


disp(sum(Joint,"all"))



Po2 = zeros(2,1);
for o2=1:2
        Po2(o2)=sum(Joint(:,:,:,o2,:), "all");
end
disp(Po2);





% Po1 in three ways

Po1_simple = zeros(2,1);
for o1=1:2
    for s1=1:2
            Po1_simple(o1) = Po1_simple(o1) + D(s1) * A(o1, s1);
    end
end
disp(Po1_simple);

Po1_mat =  A * D;
disp(Po1_mat);

% Po2 in three ways
Po2 = zeros(2,1);
for o2=1:2
        Po2(o2)=sum(Joint(:,:,:,o2,:), "all");
end
disp(Po2);

Po2_simple = zeros(2,1);
for o2=1:2
    p_s2 = 0;
    for s2=1:2
        pi_sum = 0;
        for pi=1:2
            s1_sum = 0;
            for s1=1:2
                s1_sum = s1_sum + D(s1) * B(s2,s1,pi);
            end
            pi_sum = pi_sum + E(pi) * s1_sum;
        end
        Po2_simple(o2) = Po2_simple(o2) + A(o2,s2) * pi_sum;
    end
end
disp(Po2_simple);


disp(sum(P_s1_s2_o2_pi_cnd_o1(:,:,:,:,1), "all"));
disp(sum(P_s1_s2_o2_pi_cnd_o1(:,:,:,:,2), "all"));

%p(o2|o1)
P_o2_o1 = zeros(2,1);
for o1=1:2
    for o2=1:2
        P_o2_o1(o2,o1)=sum(P_s1_s2_o2_pi_cnd_o1(:,:,o2,:,o1), "all");
    end
end
disp(P_o2_o1)

%p(s1|o1)
P_s1_o1 = zeros(2,1);
for o1=1:2
    for s1=1:2
        P_s1_o1(s1,o1)=sum(P_s1_s2_o2_pi_cnd_o1(s1,:,:,:,o1), "all");
    end
end
disp(P_s1_o1)

%p(s2|o1)
P_s2_o1 = zeros(2,1);
for o1=1:2
    for s2=1:2
        P_s2_o1(s2,o1)=sum(P_s1_s2_o2_pi_cnd_o1(:,s2,:,:,o1), "all");
    end
end
disp(P_s2_o1)

%p(pi|o1)
P_pi_o1 = zeros(2,1);
for o1=1:2
    for pi=1:2
        P_pi_o1(pi,o1)=sum(P_s1_s2_o2_pi_cnd_o1(:,:,:,pi,o1), "all");
    end
end
disp(P_pi_o1)

%p(s1,s2,o1,pi|o2)

disp(sum(P_s1_s2_o1_pi_cnd_o2(:,:,:,:,1), "all"));
disp(sum(P_s1_s2_o1_pi_cnd_o2(:,:,:,:,2), "all"));

%p(o1|o2)
P_o1_o2 = zeros(2,1);
for o2=1:2
    for o1=1:2
        P_o1_o2(o1,o2)=sum(P_s1_s2_o1_pi_cnd_o2(:,:,o1,:,o2), "all");
    end
end
disp(P_o1_o2)




Po = zeros(2,2);
for o1=1:2
    for o2=1:2
        Po(o1,o2)=sum(Joint(o1,o2,:,:,:), "all");
    end
end
disp(sum(Po,"all"))

%p(s1,s2,pi|o1,o2)
Ps1_s2_pi_o = zeros(2,2,2,2,2);
for o1 = 1:2
    for o2 = 1:2
        for s1=1:2
            for s2=1:2             
                for pi=1:2
                    Ps1_s2_pi_o(s1,s2,pi,o1,o2) = Joint(o1,o2,s1,s2,pi) / Po(o1,o2);
                end
            end
        end
    end
end
for o1=1:2
    for o2=1:2
        disp(sum(Ps1_s2_pi_o(:,:,:,o1,o2),"all"))
    end
end

Ps1_o = zeros(2,2,2);
Ps2_o = zeros(2,2,2);
Ppi_o = zeros(2,2,2);
for o2 = 1:2
    for o1 = 1:2
        Ps1_o(:,o1,o2)=sum(Ps1_s2_pi_o(:,:,:,o1,o2),[2,3]);
        Ps2_o(:,o1,o2)=sum(Ps1_s2_pi_o(:,:,:,o1,o2),[1,3]);
        Ppi_o(:,o1,o2)=sum(Ps1_s2_pi_o(:,:,:,o1,o2),[1,2]);
        disp(sum(Ps1_o(:,o1,o2),"all"))
        disp(sum(Ps2_o(:,o1,o2),"all"))
        disp(sum(Ppi_o(:,o1,o2),"all"))
    end
end
disp(Ppi_o(:,1,1))


disp('c')











% Functions
function Posterior_norm = Bayes(Likelihood, Prior, Observation)
    Posterior = Prior .* (Likelihood' * Observation);
    Posterior_norm = NORM(Posterior);
end

function A = NORM(A)
    % Normalisation of probability matrix (column elements sum to 1)
    %----------------------------------------------------------------------
    % The function goes column by column and it normalise such that 
    % elements of each colums sum to 1

    for i = 1:size(A,2)     
        for j = 1:size(A,3)
            for k = 1:size(A,4)
                for l = 1:size(A,5)
                    S = sum(A(:,i,j,k,l),1);
                    if S > 0
                        A(:,i,j,k,l) = A(:,i,j,k,l)/S;
                    else
                        A(:,i,j,k,l) = 1/size(A,1);
                    end
                end
            end
        end
    end
end